import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from qtorch.core import DataFetcher, ReturnsDataset, ModelTrainer, Strategy, Engine, Visualizer
from qtorch.model import LSTMModel, HybridSignalModel
from typing import Optional
from qtorch.strategy import SignalValidator, ThresholdStrategy
from torch.utils.data import DataLoader
from datetime import datetime
from qtorch.di.container import DIContainer
from qtorch.events.bus import EventBus

def get_akshare_fetcher():
    """延迟导入 AKShareStockFetcher 以避免循环依赖"""
    from qtorch.data import AKShareStockFetcher
    return AKShareStockFetcher

class XQBacktrader:
    def __init__(self, di_container: DIContainer, event_bus: EventBus, symbol="sz002371", start_date=None, end_date=None):
        self.di_container = di_container
        self.event_bus = event_bus
        self.fetcher = get_akshare_fetcher()()
        self.symbol = symbol
        self.end_date = pd.to_datetime(end_date or datetime.today().strftime('%Y-%m-%d'))
        default_start = (self.end_date - pd.DateOffset(years=5)).strftime('%Y-%m-%d')
        self.start_date = pd.to_datetime(start_date or default_start)
        self.model: Optional[nn.Module] = None
        self.strategy: Strategy = ThresholdStrategy(upper_threshold=0.015, lower_threshold=-0.015)
        self.validator: SignalValidator = SignalValidator(volatility_threshold=3.0, min_liquidity=1e6)
        self.use_model: bool = False
        self.model_weight: float = 0.7
        self.raw_data: Optional[pd.DataFrame] = None
        self.processed_data: Optional[pd.DataFrame] = None
        self.trainer: Optional[ModelTrainer] = None
        self.window_size: int = 30
        self.batch_size: int = 64
        self.epochs: int = 50
        self.predictions: Optional[pd.Series] = None
        self.signals: Optional[np.ndarray] = None
        self.backtest_data: Optional[pd.DataFrame] = None
        self.metrics: Optional[Dict[str, float]] = None

    def set_fetcher(self, fetcher: DataFetcher, symbol, start_date, end_date) -> None:
        self.fetcher = fetcher
        self.symbol = symbol
        self.start_date = start_date
        self.end_date = end_date

    def set_model(self, model: nn.Module, model_weight: float = 0.7, window_size: int = 30, batch_size: int = 64, epochs: int = 50):
        if not 0 <= model_weight <= 1:
            raise ValueError("Model weight must be between 0 and 1")
        self.model = model
        self.use_model = True
        self.model_weight = model_weight
        self.window_size = window_size
        self.batch_size = batch_size
        self.epochs = epochs

    def set_strategy(self, strategy: Strategy) -> None:
        self.strategy = strategy

    def __fetch_data(self, stock_code, start_date, end_date) -> None:
        self.raw_data = self.fetcher.fetch_data(stock_code, start_date, end_date)
        self.processed_data = self.fetcher.preprocess_data(self.raw_data)

    def __train_model(self) -> None:
        if not self.model:
            raise ValueError("Model must be set before training")
        self.trainer = ModelTrainer(self.model)
        self.trainer.train(self.processed_data, window_size=self.window_size, epochs=self.epochs, batch_size=self.batch_size)

    def __generate_predictions(self) -> None:
        if self.processed_data is None:
            raise ValueError("Processed data not available")
        dataset = ReturnsDataset(self.processed_data, window_size=self.window_size)
        dataloader = DataLoader(dataset, batch_size=1, shuffle=False)

        self.predictions = []
        self.model.eval()
        with torch.no_grad():
            for X, _ in dataloader:
                X = X.unsqueeze(-1).to(self.trainer.device)
                pred = self.model(X).cpu().numpy().flatten()[0]
                self.predictions.append(pred)

        self.predictions = pd.Series(self.predictions, index=self.processed_data.index[self.window_size:])

    def __generate_signals(self) -> None:
        strategy_data = self._prepare_strategy_data()
        base_signals = self._generate_base_signals(strategy_data)
        raw_signals, data_index, offset = self._combine_model_signals(base_signals)
        validated = self._validate_signal_sequence(raw_signals, offset)
        self._update_processed_data(validated)

    def _prepare_strategy_data(self) -> pd.DataFrame:
        data = self.processed_data[["open", "high", "low", "close", "volume"]].copy()
        if self.use_model:
            data["prediction"] = self.predictions.reindex(data.index, method="ffill")
        return data

    def _generate_base_signals(self, strategy_data: pd.DataFrame) -> pd.Series:
        try:
            return self.strategy.generate_signals(strategy_data)
        except KeyError as e:
            raise ValueError(f"数据缺失必要列: {e}") from e

    def _combine_model_signals(self, base_signals: pd.Series) -> tuple:
        if self.use_model:
            data_index = self.predictions.index
            offset = self.window_size
            model_part = self.model_weight * self.predictions
            raw_signals = model_part + (1 - self.model_weight) * base_signals
        else:
            data_index = self.processed_data.index
            offset = 0
            raw_signals = base_signals
        return (raw_signals, data_index, offset)

    def _validate_signal_sequence(self, raw_signals: pd.Series, offset: int) -> dict:
        validated = {"signals": [], "reasons": []}
        for i, sig in enumerate(raw_signals):
            current_data = self.processed_data.iloc[i + offset]
            is_valid, reason = self.validator.validate(
                signal=sig,
                symbol=self.symbol,
                market_data=current_data
            )
            validated["signals"].append(sig if is_valid else 0)
            validated["reasons"].append({
                "source": "hybrid" if self.use_model else "strategy",
                "status": "valid" if is_valid else reason
            })
        return validated

    def _update_processed_data(self, validated: dict):
        new_cols = pd.DataFrame({
            "signal": validated["signals"],
            "signal_source": "hybrid" if self.use_model else "strategy",
            "validation_info": validated["reasons"]
        }, index=validated["data_index"])
        self.processed_data = pd.concat([self.processed_data, new_cols], axis=1)

    def __backtest(self) -> None:
        engine = Engine(self.di_container, self.event_bus)
        self.backtest_data = engine.run([self.symbol])
        self.metrics = engine.calculate_metrics(self.backtest_data)

    def visualize(self) -> None:
        Visualizer.plot_returns(self.backtest_data)
        Visualizer.plot_trade_signals(self.backtest_data)
        Visualizer.interactive_heatmap(self.backtest_data)

    def print_results(self) -> None:
        print("\n=== 回测结果 ===")
        for k, v in self.metrics.items():
            print(f"{k}: {v:.4f}")

    def run(self) -> None:
        self.__fetch_data(self.symbol, self.start_date, self.end_date)
        if self.use_model:
            if not self.model:
                raise ValueError("Model must be set when use_model is True")
            self.__train_model()
            self.__generate_predictions()
        self.__generate_signals()
        self.__backtest()
        self.visualize()
        self.print_results()